'''
Created on Oct 30, 2009

@author: mkiyer
'''

import sys
import math
from veggie.sample.samplegroup import SampleGroup
import veggie.coverage.covdb as covdb
from plot_utr_expr import plot_utr_expr
import h5py
import numpy as np

from mirna import get_mirna_target_sites


class UTR:
    def __init__(self):
        pass
    def __str__(self):
        return '\t'.join([self.symbol, self.accession, self.chrom, self.strand, str(self.start), str(self.end), str(self.polya)])
    
    @staticmethod        
    def read_utr_copa_file(fhd):
        for line in fhd:
            if line is None:
                continue
            line = line.strip()
            if line is None:
                continue
            if line.startswith('#'):
                continue
            fields = line.split('\t')
            if len(fields) < 8:
                continue
            u = UTR()
            u.symbol = fields[0]
            u.accession = fields[1]
            u.chrom = fields[2]
            u.strand = fields[3]
            u.start = int(fields[4])
            u.end = int(fields[5])
            u.polya = int(fields[6])
            u.proximal_copas = parse_copa_field(fields[7])
            u.distal_copas = parse_copa_field(fields[8])
            yield u         

def parse_copa_field(copa_str):
    copas = [pair.split(',') for pair in copa_str.split('|')]        
    copas = [(name, float(x)) for name, x in copas]
    return copas

def get_cdf(copas):
    n = len(copas)
    rank_to_percentile = np.ones(n, dtype=np.float)
    current_score = None
    current_index = 0
    current_percentile = 0.0
    for i, copa_tuple in enumerate(copas):
        score = copa_tuple[1]
        if current_score is None:
            current_score = score
        elif score != current_score:
            current_percentile += float(i - current_index) / n
            for j in xrange(current_index, i):                    
                rank_to_percentile[j] = current_percentile
            current_index = i
            current_score = score
    return rank_to_percentile

if __name__ == '__main__':    
    proximal_threshold = 0.25
    distal_threshold = 0.25
    recurrence_threshold = 4
    covdb_file = '/archive10/bigwig/coverage/rnaseq.hdf5'
    
    benign_group = SampleGroup('benign', 'green',
                               conds=[{'app_type': 'Transcriptome',
                                       'sample_type': 'Tissue',
                                       'tissue_type': 'Breast',
                                       'diagnosis': 'Benign Tissue',
                                       'treated': 'FALSE'}])
    cancer_group = SampleGroup('tumor', 'red',
                                conds=[{'app_type': 'Transcriptome',
                                        'sample_type': 'Tissue',
                                        'tissue_type': 'Breast',
                                        'diagnosis': 'Localized Tissue',
                                        'treated': 'FALSE'}])

    benign_cl_group = SampleGroup('benign', 'pink',
                                  conds=[{'app_type': 'Transcriptome',
                                          'sample_type': 'Cell Line',
                                          'tissue_type': 'Breast',
                                          'diagnosis': 'Benign Cell Line',
                                          'treated': 'FALSE'}])    
    cancer_cl_group = SampleGroup('tumor', 'purple',
                                  conds=[{'app_type': 'Transcriptome',
                                          'sample_type': 'Cell Line',
                                          'tissue_type': 'Breast',
                                          'diagnosis': 'Localized Cell Line',
                                          'treated': 'FALSE'},
                                          {'app_type': 'Transcriptome',
                                           'sample_type': 'Cell Line',
                                           'tissue_type': 'Breast',
                                           'diagnosis': 'Metastatic Cell Line',
                                           'treated': 'FALSE'}])
    # open coverage db
    h5file = h5py.File(covdb_file, 'r')
    benign_group_cov = covdb.SampleGroupCoverage(benign_group, h5file)
    cancer_group_cov = covdb.SampleGroupCoverage(cancer_group, h5file)
    benign_cl_group_cov = covdb.SampleGroupCoverage(benign_cl_group, h5file)
    cancer_cl_group_cov = covdb.SampleGroupCoverage(cancer_cl_group, h5file)

    for u in UTR.read_utr_copa_file(open(sys.argv[1])):        
        if u.strand == '+':
            mirnas = get_mirna_target_sites(u.chrom, u.polya, u.end)
        else:
            mirnas = get_mirna_target_sites(u.chrom, u.start, u.polya)
        if len(mirnas) == 0:
            continue

        proximal_cdf = get_cdf(u.proximal_copas)
        distal_cdf = get_cdf(u.distal_copas)                
        n = int(math.ceil(proximal_threshold * len(u.proximal_copas)))        
        m = int(math.ceil(distal_threshold * len(u.distal_copas)))                
        recurrent_names = []
        
        proximal_names = set([])
        for i in xrange(n):
            if proximal_cdf[i] > proximal_threshold:
                break
            proximal_names.add(u.proximal_copas[i][0])
        distal_names = set([])
        for i in xrange(m):
            if distal_cdf[i] > distal_threshold:
                break
            distal_names.add(u.distal_copas[i][0])

        recurrent_names = proximal_names.intersection(distal_names)
        
        if len(recurrent_names) >= recurrence_threshold:
            sys.stdout.write('%s\t%s\t%s\n' % (u, ','.join(recurrent_names), ','.join(mirnas)))
            trunc_group_cov = covdb.SampleGroupCoverage(SampleGroup('truncated', 'blue', samples=recurrent_names), h5file)
            sgroups = [trunc_group_cov, benign_cl_group_cov, cancer_cl_group_cov, benign_group_cov, cancer_group_cov]
            name = '%s_%s' % (u.symbol, u.accession)
            title = '%s:%d-%d name=%s strand=%s' % \
                (u.chrom, u.start, u.end, name, u.strand)            
            outfile = '%s_%s_%d-%d_%d.png' % (name, u.chrom, u.start, u.end, u.polya) 
            plot_utr_expr(u.chrom, u.start, u.end, u.strand, u.polya, sgroups,
                          outfile=outfile, 
                          title=title,
                          shade_start=u.polya)

    h5file.close()
